package com.itheima.ai.util;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;

import java.util.List;

//有关 提示词、查询 功能
public class QueryUtil {
    //多查询扩展器（为原始提示词生成多个查询变体，提高召回文档的几率）
    public static void Expander(ChatClient chatClient) {
        MultiQueryExpander queryExpander=MultiQueryExpander.builder()
                .chatClientBuilder((ChatClient.Builder) chatClient)
                .numberOfQueries(3)
                .build();
        List<Query> queries=queryExpander.expand(new Query("谁是程序员鱼皮？"));

        //queries=
        //请介绍一下程序员鱼皮
        //请输出跟程序员鱼皮有关的背景信息
        //请介绍程序员鱼皮的个人经历
    }

    //提供了查询重写器，可以把原始提示词变得更精确和专业
    public String doQueryRewrite(ChatClient chatClient,String prompt) {
        QueryTransformer queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder((ChatClient.Builder) chatClient)
                .build();
        Query query = new Query(prompt);
        Query transformedQuery = queryTransformer.transform(query); //执行查询重写
        return transformedQuery.text();
    }
}
